#include "LU.h"

using namespace RSIM;
using namespace std;

Matrix LU::permute_copy(const Matrix &A, const Vector_<int>& piv, int j0, int j1){
        int piv_length = piv.getLength();
        Matrix X(piv_length, j1-j0+1);
        for (int i = 0; i < piv_length; i++) 
                for (int j = j0; j <= j1; j++) 
                        X(i,j-j0) = A.getElement(piv.getElement(i),j);
        return X;
}

Vector LU::permute_copy(const Vector &A,const Vector_<int> &piv){
        int piv_length = piv.getLength();
        if (piv_length != A.getLength())
                return Vector();
        Vector x(piv_length);
        for (int i = 0; i < piv_length; i++) 
                x(i) = A.getElement(piv.getElement(i));
        return x;
}

LU::LU(const Matrix &A):LU_(A), m(A.getNRow()), n(A.getNCol()), piv(A.getNRow()){

        // Use a "left-looking", dot-product, Crout/Doolittle algorithm.			
        for(int i = 0; i < m; i++)
                piv(i) = i;		
        pivsign = 1;
        double *LUrowi = 0;
        
        Vector LUcolj(m);
                // Outer loop.
                                
        for(int j = 0; j < n; j++){		
                // Make a copy of the j-th column to localize references.
                for (int i = 0; i < m; i++){
                        LUcolj[i] = LU_(i,j);
                }

                // Apply previous transformations.	
                for (int i = 0; i < m; i++) {
                        LUrowi = LU_[i];
                
                // Most of the time is spent in the following dot product.
                int kmax = min(i,j);
                double s = 0.0;
                for (int k = 0; k < kmax; k++)
                        s += LUrowi[k]*LUcolj[k];
                LUrowi[j] = LUcolj[i] -= s;
                }

                // Find pivot and exchange if necessary.
                int p = j;
                for(int i = j+1; i < m; i++){
                        if(fabs(LUcolj[i]) > fabs(LUcolj[p]))
                                p = i;
                }
                if (p != j){
                        int k=0;
                        for(k = 0; k < n; k++){
                                double t = LU_(p,k); 
                                LU_(p,k) = LU_(j,k); 
                                LU_(j,k) = t;
                        }
                        k = piv[p]; 
                        piv[p] = piv[j]; 
                        piv[j] = k;
                        pivsign = -pivsign;
                }

                // Compute multipliers.
                if((j < m) && (LU_[j][j] != 0.0)){
                        for (int i = j+1; i < m; i++){
                                LU_(i,j) =  LU_(i,j)/LU_[j][j];
                        }
                }
        }  
}

int LU::isNonsingular(){
        for(int j = 0; j < n; j++){
                if(LU_(j,j) == 0)
                        return 0;
        }
        return 1;  
}

void LU::getL(Matrix& L_Mat){
        L_Mat.setDimension(m,n);
        for(int i = 0; i < m; i++){
                for(int j = 0; j < n; j++){
                        if(i > j){
                                L_Mat(i,j) = LU_(i,j);
                        }else if(i == j){
                                L_Mat(i,j) = 1.0;
                        }else{
                                L_Mat(i,j) = 0.0;
                        }
                }
        }
}

void LU::getU(Matrix& U_Mat){
        U_Mat.setDimension(n,n);
        for(int i = 0; i < n; i++){
                for (int j = 0; j < n; j++){
                        if(i <= j){
                                U_Mat(i,j) = LU_(i,j);

                        } else {
                                U_Mat(i,j) = 0.0;
                        }
                }
        }
}

double LU::det(){
        if(m != n)
                return double(0);
        double d = double(pivsign);
        for(int j = 0; j < n; j++)
                d = d*LU_(j,j);
        return d;
}

Matrix LU::solve(const Matrix &B){

        /* Dimensions: A is mxn, X is nxk, B is mxk */

        if (B.getNRow() != m) 
                return Matrix();
        if (!isNonsingular())
                return Matrix();

// Copy right hand side with pivoting
        int nx = B.getNCol();	  
        Matrix X = permute_copy(B, piv, 0, nx-1);

        // Solve L*Y = B(piv,:)
        for (int k = 0; k < n; k++) {
                for (int i = k+1; i < n; i++) {
                        for (int j = 0; j < nx; j++) {
                                X(i,j) = X(i,j) - X(k,j)*LU_(i,k);
                        }
                }
        }
        
// Solve U*X = Y;
        for (int k = n-1; k >= 0; k--) {
                for (int j = 0; j < nx; j++) {
                        X(k,j) = X(k,j)/LU_(k,k);
                }

                for (int i = 0; i < k; i++) {
                        for (int j = 0; j < nx; j++) {
                                X(i,j) -= X(k,j)*LU_(i,k);
                        }
                }
        }
        return X;
}

Vector LU::solve(const Vector &b){

        /* Dimensions: A is mxn, X is nxk, B is mxk */
        if (b.getLength() != m) {
                return Vector();
        }
        if (!isNonsingular()) {
                return Vector();
        }
        Vector x = permute_copy(b, piv);

        // Solve L*Y = B(piv)
        for (int k = 0; k < n; k++) {
                for (int i = k+1; i < n; i++) {
                        x(i) = x(i) - x[k]*LU_(i,k);
                }
        }

        // Solve U*X = Y;
        for (int k = n-1; k >= 0; k--) {
                x[k] = x[k]/LU_(k,k);
                for (int i = 0; i < k; i++) 
                        x[i] = x[i]- x[k]*LU_(i,k);
        }
        return x;
}